import cv2
import numpy as np
import time
from gui_buttons import Buttons

#-----------------------------------UI界面按钮------------------------------------
button = Buttons()
button.add_button("person", 20, 20)
button.add_button("cell phone", 20, 100)
button.add_button("remote", 20, 180)
button.add_button("chair", 20, 260)
button.add_button("scissors", 20, 340)
colors = button.colors

#-----------------------------1、导入yolov4-tiny网络模型结构---------------------------
# 传入模型结构.cfg文件，模型权重参数.weight文件;并将模型net传入目标检测model中
net = cv2.dnn.readNet("dnn_model/yolov4-tiny.weights","dnn_model/yolov4-tiny.cfg")
model=cv2.dnn_DetectionModel(net)

#设置网络模型的输入，size表示将输入的图像缩放至320×320大小，scale 表示像素值的缩放大小。
model.setInputParams(size=(320,320),scale= 1/255)



#---------------------------------2、获取分类文本的信息--------------------------------
# 获取文本中的每一行，删除文本中的换行符、空格等，将每个分类名保存到列表classes中
classes=[]
with open("dnn_model/classes.txt","r") as f:
    for class_name in f.readlines():
        class_name = class_name.strip()
        classes.append(class_name)



#------------------------------3、摄像头开启，视频窗口与按钮界面开启-----------------------
cap = cv2.VideoCapture(0)

button_person = False
def click_button(event, x, y, flags, params):
    global button_person
    if event == cv2.EVENT_LBUTTONDOWN:
        button.button_click(x, y)

cv2.namedWindow("Frame")
cv2.setMouseCallback('Frame',click_button)


#-----------------------------------4、目标检测---------------------------------------

while True:
    #按帧读取视频
    ret,frame = cap.read()
    #UI界面按钮
    active_buttons = button.active_buttons_list()
    # 对每一帧frame图像进行目标检测  class_ids物体类别索引，scores分类置信度，bboxes检测框信息，左上角坐标(x,y)，框的宽高(w, h)
    (class_ids,scores,bboxes) = model.detect(frame)
    # 遍历所有的检测框信息                zip()函数用来并行遍历列表
    for class_id,score,bbox in zip(class_ids,scores,bboxes) :
        # 获取检测框的左上角坐标和宽高
        (x, y, w, h) = bbox
        # 获取检测框对应的分类名
        class_name = classes[class_id]
        color = colors[class_id]

        #按下哪个物体按钮，则绘制哪个物体的识别检测矩形框
        if class_name in active_buttons:
            cv2.putText(frame , class_name,(x,y-10),cv2.FONT_HERSHEY_PLAIN,2,(200,0,40),2)
            cv2.rectangle(frame, (x, y), (x + w, y + h), color, 2)

    button.display_buttons(frame)

    cv2.imshow("Frame",frame)

    #按键Esc退出
    key = cv2.waitKey(1)
    if key == 27:
        break

cap.release()
cv2.destroyAllWindows()

